from utils import print_timing, print_start_finish
import os, sys
from optparse import OptionParser

def print_error(error=None):
    print error
    sys.exit()

@print_start_finish
@print_timing
def load_all_datasets(dbname, dbtype, test, shard, recreate, jobs=1):
    import os, sys
    from utils import get_conf, pg_url
    from load_dataset import load_dataset
    conf = get_conf()
    dbuser = conf["dbuser"]
    password = conf["password"]
    portnum = conf["portnum"]
    usr = os.popen('whoami').read().strip()
    dbstring_usr = "%s://%s@:%s/postgres"%(pg_url(), usr, portnum)
    dbstring = "%s://%s:%s@localhost:%s/postgres"%(pg_url(), dbuser, password, portnum)
    dbstring_dbname = "%s://%s:%s@localhost:%s/%s"%(pg_url(), dbuser, password, portnum, dbname)
    if dbtype not in conf.sections:
        sys.exit("dbtype must be one of %s, not %s.\nWas the database %s created with init_db.py?"%(', '.join(conf.sections), dbtype, dbstring_dbname))
    from dbutils import dbuser_exists, db_exists, schema_exists, drop_schema
    from utils import create_engine_wrapper
    from sqlalchemy.orm import clear_mappers
    db_usr = create_engine_wrapper(dbstring_usr)
    db = create_engine_wrapper(dbstring)
    db_dbname = create_engine_wrapper(dbstring_dbname)
    if not db_exists(db, dbname):
        os.system("python init_db.py %s %s"%(dbname, dbtype)) and sys.exit()
    datasets = conf[dbtype].sections

    # Choose schemas depending on shard and test options.
    if shard and test:
        schemas = [dataset + "_shard_test" for dataset in datasets]
    elif shard:
        schemas = [dataset + "_shard" for dataset in datasets]
    elif test:
        schemas = [dataset + "_test" for dataset in datasets]
    else:
        schemas = datasets

    if not recreate:
        if [s for s in schemas if schema_exists(db_dbname, s, True)]:
            sys.exit("Exiting...")
    for dataset in datasets:
        clear_mappers()
        load_dataset(dataset, dbname, options.alleletype, options.jobs, options.recreate, options.shard, options.test)

usr = os.popen('whoami').read().strip()
usage = "usage: %prog [-a alleletype] [-j jobs] [-s] [-t] [-y] db dbtype"
parser = OptionParser(usage=usage)
parser.add_option("-a", "--allele-type", action="store", type="string", dest="alleletype", help="allele type")
parser.add_option("-j", "--jobs", action="store", type="int", dest="jobs", default=1, help="number of jobs to run in parallel")
parser.add_option("-s", "--sharding", action="store_true", dest="shard", help="use geno table sharding")
parser.add_option("-t", "--test", action="store_true", dest="test", default=False, help="test cases")
parser.add_option("-y", "--recreate", action="store_true", dest="recreate", help="delete schema(s) and recreate")

(options, args) = parser.parse_args()
if len(args) != 2:
    print "incorrect number of arguments"
    parser.print_help()
    sys.exit(1)
dbname = args[0]
dbtype = args[1]
print "using database %s"%(dbname)
load_all_datasets(dbname, dbtype, options.test, options.shard, options.recreate, options.jobs)
print "used database %s"%(dbname)
